import argparse
import numpy as np
import os
from models.rbfn import RBFN
from models.mlp import MLP
from models.gpr import GPR
from optimizers.de import DE
from optimizers.ga import GA
from optimizers.pso import PSO
from optimizers.rl_de import RLDE
from problems.coco_benchmark import CocoBenchmark

MODEL_DICT = {
    'rbfn': RBFN,
    'mlp': MLP,
    'gpr': GPR
}

OPTIMIZER_DICT = {
    'de': DE,
    'ga': GA,
    'pso': PSO,
    'rlde': RLDE
}

PROBLEM_DICT = {
    'sphere': CocoBenchmark.sphere,
    'rastrigin': CocoBenchmark.rastrigin,
    'rosenbrock': CocoBenchmark.rosenbrock,
    'ellipsoidal': CocoBenchmark.ellipsoidal,
    'weierstrass': CocoBenchmark.weierstrass,
    'schwefel': CocoBenchmark.schwefel
}

PROBLEM_ID_MAP = {
    'sphere': 'F1',
    'rastrigin': 'F3',
    'rosenbrock': 'F8',
    'ellipsoidal': 'F10',
    'weierstrass': 'F16',
    'schwefel': 'F20'
}

def main():
    parser = argparse.ArgumentParser(description='端到端优化框架运行脚本')
    parser.add_argument('--model', type=str, default='rbfn', choices=MODEL_DICT.keys(), help='代理模型类型')
    parser.add_argument('--optimizer', type=str, default='de', choices=OPTIMIZER_DICT.keys(), help='优化算法类型')
    parser.add_argument('--problem', type=str, default='sphere', choices=PROBLEM_DICT.keys(), help='优化问题')
    parser.add_argument('--dim', type=int, default=10, help='问题维度')
    parser.add_argument('--max_evals', type=int, default=1000, help='最大评价次数')
    args = parser.parse_args()

    # 1. 初始化问题
    func = PROBLEM_DICT[args.problem]
    dim = args.dim

    # 2. 生成初始样本
    n_init = 20 * dim
    X = np.random.uniform(-5, 5, (n_init, dim))
    # 双输出示例：y的每一列为不同目标函数
    y1 = np.array([func(x) for x in X]).reshape(-1, 1)
    y2 = np.sum(X, axis=1, keepdims=True)  # 这里以sum为第二目标，可自定义
    y = np.hstack([y1, y2])  # (N, 2)

    # 3. 初始化模型
    if args.model == 'rbfn':
        model = RBFN(input_dim=dim, output_dim=2, num_centers=20, sigma=1.0)
    elif args.model == 'mlp':
        model = MLP(input_dim=dim, hidden_dims=[64, 32], output_dim=2, lr=0.001)
    elif args.model == 'gpr':
        model = GPR(kernel='rbf', lengthscale=1.0, noise=1e-6)
    else:
        raise ValueError('未知模型类型')
    model.fit(X, y)

    # 4. 定义代理目标函数（以第一个输出为适应度）
    def surrogate_func(x):
        x = np.array(x).reshape(1, -1)
        return float(model.predict(x)[0, 0])

    # 5. 初始化优化器（只传结构参数）
    if args.optimizer == 'de':
        optimizer = DE(dim=dim, max_iter=args.max_evals)
    elif args.optimizer == 'ga':
        optimizer = GA(dim=dim, max_iter=args.max_evals)
    elif args.optimizer == 'pso':
        optimizer = PSO(dim=dim, max_iter=args.max_evals)
    elif args.optimizer == 'rlde':
        optimizer = RLDE(dim=dim, max_iter=args.max_evals)
    else:
        raise ValueError('未知优化器类型')

    # 6. 运行优化（目标函数在此传递）
    best_x, best_y = optimizer.optimize(surrogate_func)

    print(f'最优解: {best_x}')
    print(f'最优值: {best_y}')

    # 7. 对全部样本生成特征输出
    features = model.calc_activations(X)  # (N, num_centers * n_groups)
    y_pred = model.predict(X)  # (N, 2)

    # 8. 保存结果，文件名包含模型、优化器、问题编号
    os.makedirs('src/output', exist_ok=True)
    problem_id = PROBLEM_ID_MAP.get(args.problem, args.problem)
    txt_filename = f'src/output/{args.model}_{args.optimizer}_{problem_id}_log.txt'
    npy_feat_filename = f'src/output/{args.model}_{args.optimizer}_{problem_id}_features.npy'
    np.save('src/output/best_solution.npy', best_x)
    np.save(npy_feat_filename, features)
    with open(txt_filename, 'w') as f:
        f.write(f'【运行信息】\n')
        f.write(f'模型: {args.model}\n')
        f.write(f'优化器: {args.optimizer}\n')
        f.write(f'问题: {args.problem} ({problem_id})\n')
        f.write(f'维度: {args.dim}\n')
        f.write(f'最大评价次数: {args.max_evals}\n')
        f.write(f'\n【最优解】\n{best_x}\n')
        f.write(f'\n【最优值】\n{best_y}\n')
        f.write(f'\n【全部样本预测输出（双输出）】\n')
        for i, row in enumerate(y_pred):
            f.write(f'样本{i}: {row}\n')
        f.write(f'\n【全部样本中间层特征输出】\n')
        for i, row in enumerate(features):
            f.write(f'样本{i}: {row}\n')

if __name__ == '__main__':
    main() 